#include <openssl/pem.h>
#include <openssl/conf.h>
#include <openssl/x509v3.h>
#include <openssl/ssl.h>
#include <openssl/err.h>

#include <sys/poll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <string.h>

#include "mssl.h"

static int forwardData(SSL *from, SSL *to, Log *logFile);
static int isAvailable(int revents);
static int isClosed(int revents);

static void getServerName(int serverFd, char* name) {
  struct sockaddr_in serverAddr;
  int size = sizeof(serverAddr);

  if (getpeername(serverFd, (struct sockaddr*)&serverAddr, &size) < 0) {
    perror("ASSERT - s getpeername failed.");
    exit(1);
  }

  strcpy(name, inet_ntoa(serverAddr.sin_addr));
}

SSL * MSSL_handshakeWithServer(int serverFd) {
  char serverName[512];

  SSL_CTX *serverCtx;
  SSL *serverSession;

  getServerName(serverFd, serverName);

  serverCtx     = SSL_CTX_new(SSLv23_client_method());
  serverSession = SSL_new(serverCtx);


  SSL_set_connect_state(serverSession);
  SSL_set_fd(serverSession, serverFd);
  SSL_set_options(serverSession, SSL_OP_ALL);
  
  if (SSL_connect(serverSession) < 0) {
    fprintf(stderr, "Error on SSL Connect %s.\n",serverName);
    exit(1);
  }

  return serverSession;
}

SSL * MSSL_handshakeWithClient(int client, X509 *spoofedCert, Credentials *credentials) {
  SSL_CTX *clientContext = SSL_CTX_new(SSLv23_server_method());
  SSL *clientSession;

  SSL_CTX_use_certificate(clientContext, spoofedCert);
  SSL_CTX_use_PrivateKey(clientContext, credentials->leafKey);

  if (SSL_CTX_check_private_key(clientContext) == 0) {
    fprintf(stderr, "*** Assertion Failed - Generated PrivateKey Doesn't Work.\n");
    exit(1);
  }

  SSL_CTX_add_extra_chain_cert(clientContext, credentials->middleCertificate);
  SSL_CTX_set_mode(clientContext, SSL_MODE_AUTO_RETRY);

  clientSession = SSL_new(clientContext);
  SSL_set_fd(clientSession, client);

  if (SSL_accept(clientSession) == 0) {
    fprintf(stderr, "SSL Accept Failed!");
    exit(1);
  }

  return clientSession;
}


void MSSL_shuttleData(int client, SSL *clientSession, int server, SSL  *serverSession, Log *log) {
  struct pollfd fds[2] = {{client, POLLIN | POLLPRI | POLLHUP | POLLERR, 0},
			  {server, POLLIN | POLLPRI | POLLHUP | POLLERR, 0}};

  for (;;) {
    if (poll(fds, 2, -1) < 0) return;

    if (isAvailable(fds[0].revents)) if (forwardData(clientSession, serverSession, log) != 0) 
	{ 
char serverName[50]; 
getServerName(server, serverName);
fprintf(stderr, "forwardData client to server %s failed! \n",serverName);
return;}
    if (isAvailable(fds[1].revents)) if (forwardData(serverSession, clientSession, log) != 0) 
	{
char serverName[50]; 
getServerName(server, serverName);
fprintf(stderr, "forwardData server %s to client failed! \n",serverName);return;}
    if (isClosed(fds[0].revents)) return;
    if (isClosed(fds[0].revents)) return;
  }
}

static int isAvailable(int revents) {
  return revents & POLLIN || revents & POLLPRI;
}

static int isClosed(int revents) {
  return revents & POLLERR || revents & POLLHUP;
}

static int forwardData(SSL *from, SSL *to, Log *log) {
  char buf[4096];
  int bytesRead; 
  int bytesWritten;

  do {
    if ((bytesRead    = SSL_read(from, buf, sizeof(buf))) <= 0)    
{    fprintf(stderr, "SSL_read Failed! bytesRead %d\n",bytesRead);
 return -1;}
    if ((bytesWritten = SSL_write(to, buf, bytesRead)) < bytesRead) 
{	    fprintf(stderr, "SSL_write Failed!\n");
return -1;}

    LOG_log(log, buf, bytesRead);
  } while (SSL_pending(from));

  LOG_flush(log);

  return 0;
}
